#!/usr/bin/env python3
"""
肤色融合功能使用示例
展示如何在实际项目中使用肤色融合功能
"""

import sys
import os
import numpy as np
from PIL import Image
import cv2

# 添加当前目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from skin_color_fusion import skin_color_fusion_pipeline, SkinColorFusionNode

def load_and_preprocess_image(image_path: str, target_size: tuple = (768, 1440)) -> np.ndarray:
    """加载并预处理图像"""
    if not os.path.exists(image_path):
        raise FileNotFoundError(f"图像文件不存在: {image_path}")
    
    img = Image.open(image_path)
    img = img.convert('RGB')
    img = img.resize(target_size, Image.Resampling.LANCZOS)
    return np.array(img)

def create_skin_mask_from_image(image: np.ndarray, skin_color_range: tuple = None) -> np.ndarray:
    """
    从图像中自动创建肤色mask
    这是一个简化的示例，实际项目中可能需要更复杂的算法
    """
    # 转换到HSV色彩空间
    hsv = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
    
    # 定义肤色范围（HSV）
    if skin_color_range is None:
        # 默认肤色范围
        lower_skin = np.array([0, 20, 70], dtype=np.uint8)
        upper_skin = np.array([20, 255, 255], dtype=np.uint8)
    else:
        lower_skin, upper_skin = skin_color_range
    
    # 创建肤色mask
    mask = cv2.inRange(hsv, lower_skin, upper_skin)
    
    # 形态学操作，去除噪声
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
    
    return mask

def process_live2d_skin_fusion(
    reference_image_path: str,
    white_model_path: str,
    output_path: str,
    reference_skin_mask_path: str = None,
    white_model_mask_path: str = None
):
    """
    处理Live2D肤色融合的完整流程
    
    Args:
        reference_image_path: 参考图像路径
        white_model_path: 白模图像路径
        output_path: 输出图像路径
        reference_skin_mask_path: 参考图肤色mask路径（可选，会自动生成）
        white_model_mask_path: 白模mask路径（可选，会自动生成）
    """
    print(f"加载参考图像: {reference_image_path}")
    reference_image = load_and_preprocess_image(reference_image_path)
    
    print(f"加载白模图像: {white_model_path}")
    white_model = load_and_preprocess_image(white_model_path)
    
    # 处理参考图肤色mask
    if reference_skin_mask_path and os.path.exists(reference_skin_mask_path):
        print(f"加载参考图肤色mask: {reference_skin_mask_path}")
        reference_skin_mask = load_and_preprocess_image(reference_skin_mask_path)
        reference_skin_mask = cv2.cvtColor(reference_skin_mask, cv2.COLOR_RGB2GRAY)
    else:
        print("自动生成参考图肤色mask...")
        reference_skin_mask = create_skin_mask_from_image(reference_image)
    
    # 处理白模mask
    if white_model_mask_path and os.path.exists(white_model_mask_path):
        print(f"加载白模mask: {white_model_mask_path}")
        white_model_mask = load_and_preprocess_image(white_model_mask_path)
        white_model_mask = cv2.cvtColor(white_model_mask, cv2.COLOR_RGB2GRAY)
    else:
        print("自动生成白模mask...")
        # 对于白模，我们假设整个图像区域都需要处理
        white_model_mask = np.ones(white_model.shape[:2], dtype=np.uint8) * 255
    
    print("执行肤色融合...")
    fused_image, fusion_mask = skin_color_fusion_pipeline(
        reference_image, reference_skin_mask, white_model, white_model_mask,
        blend_strength=0.8,
        color_adjustment=0.3,
        texture_strength=0.2,
        feather_radius=5
    )
    
    print(f"保存融合结果: {output_path}")
    Image.fromarray(fused_image).save(output_path)
    
    # 保存融合mask
    mask_output_path = output_path.replace('.png', '_mask.png')
    Image.fromarray(fusion_mask).save(mask_output_path)
    
    print("处理完成！")
    print(f"- 融合结果: {output_path}")
    print(f"- 融合mask: {mask_output_path}")

def create_comfyui_workflow_example():
    """创建ComfyUI工作流示例"""
    workflow = {
        "nodes": [
            {
                "id": 1,
                "type": "LoadImage",
                "name": "加载参考图像",
                "inputs": {"image": "reference_image.png"}
            },
            {
                "id": 2,
                "type": "LoadImage", 
                "name": "加载参考图mask",
                "inputs": {"image": "reference_mask.png"}
            },
            {
                "id": 3,
                "type": "LoadImage",
                "name": "加载白模图像", 
                "inputs": {"image": "white_model.png"}
            },
            {
                "id": 4,
                "type": "LoadImage",
                "name": "加载白模mask",
                "inputs": {"image": "white_model_mask.png"}
            },
            {
                "id": 5,
                "type": "SkinColorFusion",
                "name": "肤色融合",
                "inputs": {
                    "reference_image": [1, 0],
                    "reference_skin_mask": [2, 0],
                    "white_model": [3, 0],
                    "white_model_mask": [4, 0],
                    "blend_strength": 0.8,
                    "color_adjustment": 0.3,
                    "texture_strength": 0.2,
                    "feather_radius": 5
                }
            },
            {
                "id": 6,
                "type": "SaveImage",
                "name": "保存融合结果",
                "inputs": {"images": [5, 0]}
            }
        ]
    }
    
    return workflow

if __name__ == "__main__":
    # 示例用法
    print("肤色融合功能使用示例")
    print("=" * 50)
    
    # 如果有实际的图像文件，可以取消注释下面的代码进行测试
    """
    try:
        process_live2d_skin_fusion(
            reference_image_path="path/to/reference_image.png",
            white_model_path="path/to/white_model.png", 
            output_path="output/fused_result.png",
            reference_skin_mask_path="path/to/reference_mask.png",
            white_model_mask_path="path/to/white_model_mask.png"
        )
    except Exception as e:
        print(f"处理失败: {e}")
    """
    
    print("\nComfyUI工作流示例:")
    workflow = create_comfyui_workflow_example()
    print(workflow)
